# ensemble.py

import numpy as np
from config.config import *
import joblib
import pandas as pd
class AverageEnsemble:
    """
    最简单的“概率平均”集成：
    pipelines: list of 已 fit 好的 Pipeline，每个 pipeline 必须有 predict_proba() 方法。
    threshold: 二分类阈值，默认 0.5，predict() 直接根据 avg_prob >= threshold 判正类。
    """
    def __init__(self, pipelines):
        if len(pipelines) == 0:
            raise ValueError("pipelines 列表不能为空")
        self.pipelines = pipelines
        path = opj(weight_result_path,"ensemble_model")    
        self.imp = joblib.load(opj(path,'imputer_median.pkl'))
        self.scaler = joblib.load(opj(path,'scaler_standard.pkl'))

    def predict_proba(self, X):
        probas = []
        for mdl in self.pipelines:
            p_i = mdl.predict_proba(X)[:, 1].reshape(-1, 1)
            probas.append(p_i)
        probas_mat = np.hstack(probas)
        avg_pos = np.mean(probas_mat, axis=1)
        avg_neg = 1.0 - avg_pos
        return np.vstack([avg_neg, avg_pos]).T


    def predict_proba_by_dif(self,X_test):
        X_test_imputed = pd.DataFrame(
            self.imp.transform(X_test),
            columns=X_test.columns,
            index=X_test.index
        )
        # 归一化分别处理
        X_test_scalered = pd.DataFrame(
            self.scaler.transform(X_test_imputed),
            columns=X_test.columns,
            index=X_test.index
        )
        proba_list = []
        for name in self.pipelines.keys():
            X_process = X_test_imputed
            if name in need_standard_set:
                X_process = X_test_scalered
            # 多个模型pipeline的预测平均
            for mdl in self.pipelines[name]:
                proba_list.append(mdl.predict_proba(X_process)[:, 1].reshape(-1, 1))
        probas_mat = np.hstack(proba_list)
        avg_pos = np.mean(probas_mat, axis=1)
        avg_neg = 1.0 - avg_pos
        return np.vstack([avg_neg, avg_pos]).T